package com.zxy.zxyrpc.registry;

import cn.hutool.core.collection.ConcurrentHashSet;
import com.zxy.zxyrpc.config.RegistryConfig;
import com.zxy.zxyrpc.exception.RpcException;
import com.zxy.zxyrpc.model.ErrorCode;
import com.zxy.zxyrpc.model.ServiceMetaInfo;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * ZooKeeper 注册中心
 *
 * @author zxy
 * @since 2025/5/18
 **/
public class ZookeeperRegistry implements Registry {
    /**
     * 根节点（业务前缀，basePath 结尾不能带 /）
     */
    private static final String ZK_ROOT_PATH = "/rpc";

    /**
     * 本机注册节点 key 集合（用于维护心跳续期）
     */
    private final Set<String> localRegisterNodeKeySet = new HashSet<>();
    /**
     * 注册服务缓存
     */
    private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();
    /**
     * 监听的 key 集合（ConcurrentHashSet 防止并发冲突）
     */
    private final Set<String> watchingKeySet = new ConcurrentHashSet<>();
    private CuratorFramework client;
    private ServiceDiscovery<ServiceMetaInfo> serviceDiscovery;

    @Override
    public void init(RegistryConfig registryConfig) {
        // String address = ZK_ADDRESS + registryConfig.getAddress();
        // 构建 client 实例
        client = CuratorFrameworkFactory
                .builder()
                .connectString(registryConfig.getAddress())
                .retryPolicy(new ExponentialBackoffRetry(Math.toIntExact(registryConfig.getTimeout()), 3))
                .build();
        // 构建 serviceDiscovery 实例
        serviceDiscovery = ServiceDiscoveryBuilder
                .builder(ServiceMetaInfo.class)
                .client(client)
                .basePath(ZK_ROOT_PATH)
                .serializer(new JsonInstanceSerializer<>(ServiceMetaInfo.class))
                .build();

        // 启动 client 和 serviceDiscovery
        try {
            client.start();
            serviceDiscovery.start();
        } catch (Exception e) {
            throw new RpcException(ErrorCode.SYSTEM_ERROR, "zk 注册中心启动失败");
        }
        // heartBeat();
    }

    @Override
    public void registry(ServiceMetaInfo serviceMetaInfo) throws Exception {
        serviceDiscovery.registerService(buildServiceInstance(serviceMetaInfo));
        String registerKey = ZK_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();
        // 添加节点信息到本地缓存
        localRegisterNodeKeySet.add(registerKey);
    }

    @Override
    public void unRegistry(ServiceMetaInfo serviceMetaInfo) {
        try {
            serviceDiscovery.unregisterService(buildServiceInstance(serviceMetaInfo));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        String registerKey = ZK_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();
        localRegisterNodeKeySet.remove(registerKey);
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        // 先走服务缓存
        List<ServiceMetaInfo> serviceMetaInfos = registryServiceCache.readCache(serviceKey);
        if (serviceMetaInfos != null)
            return serviceMetaInfos;
        // String queryKey = ZK_ROOT_PATH + "/" + serviceKey;
        try {
            // 查询服务列表信息
            Collection<ServiceInstance<ServiceMetaInfo>> serviceInstanceList = serviceDiscovery.queryForInstances(serviceKey);
            // 解析服务列表信息
            List<ServiceMetaInfo> serviceMetaInfoList = serviceInstanceList.stream()
                    .map(ServiceInstance::getPayload)
                    .collect(Collectors.toList());

            // 写入服务缓存
            registryServiceCache.writeCache(serviceKey, serviceMetaInfoList);
            return serviceMetaInfoList;
        } catch (Exception e) {
            throw new RpcException(ErrorCode.OPERATION_ERROR, "服务列表获取失败");
        }

    }

    @Override
    public void destroy() {
        System.out.println("当前服务节点下线");
        // 下线节点（这一步可以不做，因为都是临时节点，服务下线，自然都被删除掉了）
        for (String key : localRegisterNodeKeySet) {
            try {
                client.delete().guaranteed().forPath(key);
            } catch (Exception e) {
                throw new RpcException(ErrorCode.OPERATION_ERROR, key + "节点下线失败");
            }
        }
        // 释放资源
        if (client != null)
            client.close();
        /*
        可能client会帮我们关闭
        if (serviceDiscovery != null) {
            try {
                serviceDiscovery.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }*/
    }

    /**
     * zookeeper 自带心跳检测
     */
    @Override
    public void heartBeat() {
        //     不需要心跳机制，建立了临时节点，如果服务器故障，则临时节点自动丢失
    }

    // todo 可优化
    @Override
    public void watch(String serviceKey, String serviceNodeKey) {
        String watchKey = ZK_ROOT_PATH + "/" + serviceNodeKey;
        boolean newWatch = watchingKeySet.add(watchKey);
        if (newWatch) {
            CuratorCache curatorCache = CuratorCache.build(client, watchKey);
            curatorCache.start();
            curatorCache.listenable().addListener(
                    CuratorCacheListener.builder()
                            .forDeletes(childData -> registryServiceCache.clearCache())
                            .forChanges(((oldNode, node) -> registryServiceCache.clearCache()))
                            .build()
            );
        }
    }

    /**
     * 创建服务请求实例
     */
    private ServiceInstance<ServiceMetaInfo> buildServiceInstance(ServiceMetaInfo serviceMetaInfo) {
        String serviceAddress = serviceMetaInfo.getServiceHost() + ":" + serviceMetaInfo.getServicePort();
        try {
            return ServiceInstance
                    .<ServiceMetaInfo>builder()
                    .id(serviceAddress)
                    .name(serviceMetaInfo.getServiceKey())
                    .address(serviceAddress)
                    .payload(serviceMetaInfo)
                    .build();
        } catch (Exception e) {
            throw new RpcException(ErrorCode.SYSTEM_ERROR,"创建zookeeper实例失败");
        }
    }

}
